# -*- coding: utf-8 -*-
# @author: HRUN
"""
数据库管理模块 - 处理多线程环境下的数据库连接
"""

import threading
from django.db import connection, close_old_connections, reset_queries
from django.db.utils import DatabaseError
import functools

# 线程本地存储
_thread_local = threading.local()

def get_thread_local_db():
    """获取线程本地的数据库连接"""
    if not hasattr(_thread_local, 'db_connection'):
        # 关闭旧的连接
        close_old_connections()
        # 重置查询日志
        reset_queries()
        # 标记当前线程已初始化
        _thread_local.db_connection = True
    return _thread_local.db_connection

def thread_safe_db_operation(func):
    """装饰器：确保数据库操作在正确的线程中执行"""
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        try:
            # 确保数据库连接正确
            get_thread_local_db()

            # 执行数据库操作
            result = func(*args, **kwargs)

            # 关闭连接
            connection.close()
            
            return result
        except DatabaseError as e:
            # print(f"数据库操作失败: {e}")
            # 尝试重新连接
            try:
                close_old_connections()
                reset_queries()
                result = func(*args, **kwargs)
                connection.close()
                return result
            except Exception as e2:
                # print(f"重试数据库操作失败: {e2}")
                raise
        except Exception as e:
            print(f"数据库操作异常: {e}")
            raise
    
    return wrapper

def safe_get_env_config(env_id, debug=True):
    """线程安全的获取环境配置"""
    @thread_safe_db_operation
    def _get_env_config():
        from projects.models import TestEnv
        
        env = TestEnv.objects.get(pk=env_id)
        var = {**env.global_variable, **env.debug_global_variable} if debug else env.global_variable
        # 环境变量
        ENV = {
            **var,  # 因为debug模式下var中有临时的调试变量，解包时不能写到下面，否则会覆盖真正的host和headers
            'host': env.host,
            'headers': env.headers,
        }
        config = {
            'ENV': ENV,
            'DB': env.db,
            'global_func': env.global_func
        }
        return config
    
    try:
        return _get_env_config()
    except Exception as e:
        # print(f"获取环境配置失败 (env_id={env_id}): {e}")
        # 返回默认配置
        return {
            'ENV': {
                'host': 'http://localhost',
                'headers': {},
                'global_func': b''
            },
            'DB': {},
            'global_func': b''
        }

def safe_save_log(report_id, level, message, **kwargs):
    """线程安全的保存日志"""
    @thread_safe_db_operation
    def _save_log():
        from performance.models import PerformanceTestLog
        from django.utils import timezone
        
        # 获取字段，并做类型规整
        user_count = kwargs.get('user_count')
        request_name = kwargs.get('request_name')
        response_time = kwargs.get('response_time')
        
        # user_count 类型检查和转换
        if user_count is not None:
            try:
                # 确保是数字类型
                if isinstance(user_count, (int, float)):
                    user_count = int(user_count)
                elif isinstance(user_count, str) and user_count.strip():
                    user_count = int(float(user_count))
                else:
                    user_count = None
            except (ValueError, TypeError):
                user_count = None

        # 保护 message，避免作用域冲突与 None 值
        msg = message
        if msg is None:
            msg = ""

        # 规整 request_name
        if request_name is not None and not isinstance(request_name, str):
            try:
                request_name = str(request_name)
            except Exception:
                request_name = None

        # 规整 response_time
        if response_time is not None and not isinstance(response_time, (int, float)):
            try:
                response_time = float(response_time)
            except (ValueError, TypeError):
                response_time = None
        
        # 创建日志记录
        log = PerformanceTestLog(
            report_id=report_id,
            timestamp=timezone.now(),
            level=level,
            message=msg,
            source='locust',
            user_count=user_count,
            request_name=request_name,
            response_time=response_time,
            exception=kwargs.get('exception')
        )
        
        # 在保存之前验证模型
        try:
            log.full_clean()
        except Exception as e:
            # 如果验证失败，尝试修复数据
            if log.user_count is not None and not isinstance(log.user_count, int):
                try:
                    log.user_count = int(log.user_count) if log.user_count is not None else None
                except (ValueError, TypeError):
                    log.user_count = None
        
        # 最后的保护措施：确保所有字段都是正确的类型
        try:
            # 确保 user_count 是整数或 None
            if log.user_count is not None:
                if not isinstance(log.user_count, int):
                    log.user_count = int(log.user_count) if log.user_count is not None else None
        except (ValueError, TypeError):
            log.user_count = None
        
        # 尝试保存，如果失败则记录详细错误
        try:
            log.save()
            return log
        except Exception as save_error:
            # 如果保存失败，尝试使用更安全的方式
            try:
                # 重新创建对象，确保类型正确
                safe_log = PerformanceTestLog(
                    report_id=log.report_id,
                    timestamp=log.timestamp,
                    level=log.level,
                    message=log.message or "",
                    source=log.source,
                    user_count=None,  # 强制设置为 None
                    request_name=log.request_name,
                    response_time=log.response_time,
                    exception=log.exception
                )
                safe_log.save()
                return safe_log
            except Exception:
                return None
    
    try:
        return _save_log()
    except Exception as e:
        print(f"保存日志失败: {e}")
        return None 